import numpy as np

class NMS:
    """
    筛选出子任务所需的指定类型的对象数据
    """

    def __init__(self, keys, NMS_threshold = 0.6,NMS_metric = 'iou'):
        self.keys = keys
        self.match_threshold = NMS_threshold
        self.match_metric = NMS_metric

    def __call__(self, data: dict):
        
        batch_bboxes = data[self.keys["in_det"]]["boxes"]
        target_bboxes = []
        for bboxes in batch_bboxes:
            bboxes = self.nms(bboxes)
            target_bboxes.append(bboxes)
        data[self.keys["out_det"]] = {
            "boxes": target_bboxes,
        }

        return data
    
    def nms(self, dets):
        """ Apply NMS to avoid detecting too many overlapping bounding boxes.
            Args:
                dets: shape [N, 5], [score, x1, y1, x2, y2]
                match_metric: 'iou' or 'ios'
                match_threshold: overlap thresh for match metric.
        """
        if dets.shape[0] == 0:
            return dets
        scores = dets[:, 1]
        x1 = dets[:, 2]
        y1 = dets[:, 3]
        x2 = dets[:, 4]
        y2 = dets[:, 5]
        areas = (x2 - x1 + 1) * (y2 - y1 + 1)
        order = scores.argsort()[::-1]

        ndets = dets.shape[0]
        suppressed = np.zeros((ndets), dtype=np.int32)

        for _i in range(ndets):
            i = order[_i]
            if suppressed[i] == 1:
                continue
            ix1 = x1[i]
            iy1 = y1[i]
            ix2 = x2[i]
            iy2 = y2[i]
            iarea = areas[i]
            for _j in range(_i + 1, ndets):
                j = order[_j]
                if suppressed[j] == 1:
                    continue
                xx1 = max(ix1, x1[j])
                yy1 = max(iy1, y1[j])
                xx2 = min(ix2, x2[j])
                yy2 = min(iy2, y2[j])
                w = max(0.0, xx2 - xx1 + 1)
                h = max(0.0, yy2 - yy1 + 1)
                inter = w * h
                if self.match_metric == 'iou':
                    union = iarea + areas[j] - inter
                    match_value = inter / union
                elif self.match_metric == 'ios':
                    smaller = min(iarea, areas[j])
                    match_value = inter / smaller
                else:
                    raise ValueError()
                if match_value >= self.match_threshold:
                    suppressed[j] = 1
        keep = np.where(suppressed == 0)[0]
        dets = dets[keep, :]
        return dets
